#!/usr/bin/env python3
"""
实时IMU振动监控模块 - 用于Web监控系统
通过蓝牙连接IMU设备，提供实时振动数据
"""
import asyncio
import time
import threading
import numpy as np
from queue import Queue
from datetime import datetime
import logging
from bleak import BleakClient

# IMU 特征 UUID
IMU_CONTROL_CHARACTERISTIC_UUID = "0000ae01-0000-1000-8000-00805f9b34fb"
IMU_DATA_CHARACTERISTIC_UUID = "0000ae02-0000-1000-8000-00805f9b34fb"

# 数据解析比例因子
scaleAccel = 0.00478515625      # 加速度 [-16g~+16g]
scaleQuat = 0.000030517578125   # 四元数 [-1~+1]
scaleAngle = 0.0054931640625    # 角度 [-180~+180]
scaleAngleSpeed = 0.06103515625 # 角速度 [-2000~+2000]

class RealtimeIMUMonitor:
    """实时IMU振动监控类"""

    def __init__(self, mac_address="26:EF:60:96:87:3A"):
        """
        初始化IMU监控

        Args:
            mac_address: 蓝牙设备MAC地址
        """
        self.mac_address = mac_address
        self.running = False
        self.data_queue = Queue(maxsize=100)
        self.thread = None
        self.logger = logging.getLogger(__name__)

        # 蓝牙客户端
        self.client = None
        self.loop = None

        # 数据解析标志
        self.parse_enabled = False

        # 线程退出标志和事件
        self._stop_event = threading.Event()
        self._thread_stopped = threading.Event()
        self._cleanup_done = False

    def parse_imu(self, buf):
        """解析 IMU 数据"""
        try:
            self.logger.debug(f"开始解析IMU数据，第一个字节: 0x{buf[0]:02x}")

            imu_data = {
                'ax': 0.0, 'ay': 0.0, 'az': 0.0,
                'gx': 0.0, 'gy': 0.0, 'gz': 0.0,
                'roll': 0.0, 'pitch': 0.0, 'yaw': 0.0,
                'timestamp': datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f")[:-3]
            }

            if buf[0] == 0x11:
                ctl = (buf[2] << 8) | buf[1]
                L = 7  # 从第7字节开始解析数据

                # 解析加速度数据 (ctl & 0x0001)
                if (ctl & 0x0001) != 0:
                    for i in range(3):
                        value = np.short((np.short(buf[L+1]) << 8) | buf[L]) * scaleAccel
                        if i == 0:
                            imu_data['ax'] = float(value)
                        elif i == 1:
                            imu_data['ay'] = float(value)
                        else:
                            imu_data['az'] = float(value)
                        L += 2

                # 解析角速度数据 (ctl & 0x0002)
                if (ctl & 0x0002) != 0:
                    for i in range(3):
                        value = np.short((np.short(buf[L+1]) << 8) | buf[L]) * scaleAngleSpeed
                        if i == 0:
                            imu_data['gx'] = float(value)
                        elif i == 1:
                            imu_data['gy'] = float(value)
                        else:
                            imu_data['gz'] = float(value)
                        L += 2

                # 解析角度数据 (ctl & 0x0004)
                if (ctl & 0x0004) != 0:
                    for i in range(3):
                        value = np.short((np.short(buf[L+1]) << 8) | buf[L]) * scaleAngle
                        if i == 0:
                            imu_data['roll'] = float(value)
                        elif i == 1:
                            imu_data['pitch'] = float(value)
                        else:
                            imu_data['yaw'] = float(value)
                        L += 2

                # 计算振动幅度
                magnitude = np.sqrt(imu_data['ax']**2 + imu_data['ay']**2 + imu_data['az']**2)
                imu_data['magnitude'] = float(magnitude)

                self.logger.debug(f"解析成功: magnitude={magnitude:.3f}g")

                # 放入队列
                if not self.data_queue.full():
                    self.data_queue.put(imu_data)
                else:
                    # 队列满时丢弃旧数据
                    try:
                        self.data_queue.get_nowait()
                        self.data_queue.put(imu_data)
                    except:
                        pass
            else:
                self.logger.warning(f"数据包格式错误，第一个字节应为0x11，实际为0x{buf[0]:02x}")

        except Exception as e:
            self.logger.error(f"IMU数据解析错误: {e}")
            import traceback
            traceback.print_exc()

    def notification_handler(self, sender, data):
        """处理从 IMU 设备接收到的数据通知"""
        if self.parse_enabled and self.running:
            self.parse_imu(data)

    async def configure_device_async(self, client):
        """配置 IMU 设备（异步）"""
        try:
            self.logger.info("配置IMU设备...")

            # 打印所有服务和特征值用于调试
            self.logger.info("列出所有可用的服务和特征值:")
            for service in client.services:
                self.logger.info(f"  服务: {service.uuid}")
                for char in service.characteristics:
                    self.logger.info(f"    特征: {char.uuid} (handle={char.handle}, properties={char.properties})")

            # 查找控制特征和数据特征 - 明确使用服务 0000ae30 下的特征
            TARGET_SERVICE_UUID = "0000ae30-0000-1000-8000-00805f9b34fb"
            control_char = None
            data_char = None

            for service in client.services:
                # 只从目标服务中查找特征
                if service.uuid.lower() == TARGET_SERVICE_UUID.lower():
                    self.logger.info(f"🎯 使用目标服务: {service.uuid}")
                    for char in service.characteristics:
                        if char.uuid.lower() == IMU_CONTROL_CHARACTERISTIC_UUID.lower():
                            control_char = char
                            self.logger.info(f"✅ 找到控制特征: {char.uuid} (handle={char.handle})")
                        elif char.uuid.lower() == IMU_DATA_CHARACTERISTIC_UUID.lower():
                            data_char = char
                            self.logger.info(f"✅ 找到数据特征: {char.uuid} (handle={char.handle})")
                    break  # 找到目标服务后退出

            if not control_char or not data_char:
                self.logger.error(f"未找到必需的特征值! control={control_char}, data={data_char}")
                self.logger.error(f"期望的UUID: control={IMU_CONTROL_CHARACTERISTIC_UUID}, data={IMU_DATA_CHARACTERISTIC_UUID}")
                return None

            self.logger.info("发送配置命令...")

            # 发送配置命令 - 使用handle
            await client.write_gatt_char(control_char.handle, b')')
            await asyncio.sleep(0.1)

            await client.write_gatt_char(control_char.handle, bytes([0x46]))
            await asyncio.sleep(0.1)

            # 参数设置
            params = bytearray([0x00] * 11)
            params[0] = 0x12
            params[1] = 5
            params[2] = 255
            params[3] = 0
            params[4] = 0
            params[5] = 60      # 60Hz上报频率
            params[6] = 1
            params[7] = 3
            params[8] = 5
            params[9] = 0xFF
            params[10] = 0x0F

            await client.write_gatt_char(control_char.handle, bytes(params))
            await asyncio.sleep(0.1)

            await client.write_gatt_char(control_char.handle, bytes([0x19]))
            await asyncio.sleep(0.1)

            # 先尝试停止可能存在的旧通知（处理上次异常退出的情况）
            try:
                self.logger.info("尝试清理可能存在的旧通知订阅...")
                await client.stop_notify(data_char.handle)
                await asyncio.sleep(0.3)
                self.logger.info("旧通知已清理")
            except Exception as e:
                self.logger.debug(f"清理旧通知失败（可能不存在）: {e}")

            # 启用通知 - 使用handle
            self.logger.info("启用数据通知...")
            await client.start_notify(data_char.handle, self.notification_handler)

            # 等待一下，确保通知已经激活
            await asyncio.sleep(0.5)

            self.logger.info("✅ IMU设备配置完成!")
            return data_char.handle  # 返回数据特征句柄，用于后续停止通知

        except Exception as e:
            self.logger.error(f"配置IMU设备失败: {e}")
            import traceback
            traceback.print_exc()
            return None

    async def run_async(self):
        """异步运行循环"""
        client = None
        data_char_handle = None
        try:
            self.logger.info(f"连接IMU设备: {self.mac_address}")

            # 先扫描设备
            self.logger.info("扫描蓝牙设备...")
            from bleak import BleakScanner
            devices = await BleakScanner.discover(timeout=5.0)

            found = False
            for device in devices:
                if device.address.upper() == self.mac_address.upper():
                    found = True
                    self.logger.info(f"✅ 找到设备: {device.name} ({device.address})")
                    break

            if not found:
                self.logger.error(f"❌ 未找到设备 {self.mac_address}")
                self.logger.info(f"已扫描到的设备: {[f'{d.name}({d.address})' for d in devices[:5]]}")
                return

            client = BleakClient(self.mac_address, timeout=30.0)
            await client.connect()
            self.client = client

            if client.is_connected:
                self.logger.info("✅ IMU设备连接成功!")

                await asyncio.sleep(2)  # 等待服务发现

                # 配置设备并获取数据特征句柄
                config_result = await self.configure_device_async(client)
                if config_result:
                    data_char_handle = config_result  # 保存数据特征句柄
                    self.parse_enabled = True
                    self.logger.info("开始接收IMU数据...")

                    # 保持连接，同时监听停止事件
                    while self.running and not self._stop_event.is_set():
                        await asyncio.sleep(0.1)

                    # 收到停止信号，开始清理
                    pass
                else:
                    # 配置失败
                    pass
            else:
                self.logger.error("IMU设备未正确连接")

        except Exception as e:
            self.logger.error(f"IMU连接错误: {e}")
            import traceback
            traceback.print_exc()
        finally:
            # 完全清理蓝牙资源
            self.parse_enabled = False

            # 停止通知 - 关键清理步骤，必须执行
            if client and client.is_connected and data_char_handle:
                try:
                    await client.stop_notify(data_char_handle)
                    await asyncio.sleep(0.2)  # 等待停止完成
                except Exception as e:
                    pass

            # 断开蓝牙连接 - 关键清理步骤，必须执行
            if client and client.is_connected:
                try:
                    await client.disconnect()
                    # 额外等待确保断开完成 - 这很重要！
                    await asyncio.sleep(1.0)  # 增加等待时间，确保蓝牙完全断开
                except Exception as e:
                    pass

            self.client = None
            self._cleanup_done = True
            self._thread_stopped.set()

    def _run_loop(self):
        """在独立线程中运行事件循环"""
        self.loop = asyncio.new_event_loop()
        asyncio.set_event_loop(self.loop)

        try:
            self.loop.run_until_complete(self.run_async())
        except Exception as e:
            self.logger.error(f"IMU事件循环错误: {e}")
        finally:
            # 彻底清理事件循环
            try:
                # 1. 取消所有待处理的任务
                pending = asyncio.all_tasks(self.loop)
                for task in pending:
                    task.cancel()

                # 2. 等待所有任务取消 - 关键步骤
                if pending:
                    self.loop.run_until_complete(asyncio.gather(*pending, return_exceptions=True))
                    # 等待任务完全取消
                    time.sleep(0.2)

                # 3. 关闭所有异步生成器
                self.loop.run_until_complete(self.loop.shutdown_asyncgens())

                # 4. 关闭所有执行器
                self.loop.run_until_complete(self.loop.shutdown_default_executor())

            except Exception as e:
                pass
            finally:
                # 5. 关闭事件循环 - 必须关闭，否则下次无法创建新循环
                try:
                    self.loop.close()
                    # 等待确保关闭完成
                    time.sleep(0.2)
                except Exception as e:
                    pass

                self.loop = None
                self._thread_stopped.set()

    def start(self):
        """启动IMU监控"""
        if self.running:
            self.logger.warning("IMU监控已在运行中")
            return

        # 重置退出标志和事件
        self._stop_event.clear()
        self._thread_stopped.clear()
        self._cleanup_done = False

        # 清空数据队列，防止旧数据
        while not self.data_queue.empty():
            try:
                self.data_queue.get_nowait()
            except:
                break

        self.running = True
        self.thread = threading.Thread(target=self._run_loop, daemon=False)  # 改为非daemon线程
        self.thread.start()
        self.logger.info("IMU监控已启动")

    def stop(self):
        """停止IMU监控"""
        if not self.running:
            return

        self.running = False
        self.parse_enabled = False

        # 设置停止事件
        self._stop_event.set()

        # 等待线程完全退出（最多15秒） - 给蓝牙断开足够时间
        if self.thread and self.thread.is_alive():
            success = self._thread_stopped.wait(timeout=15)
            # 再等待线程对象
            self.thread.join(timeout=3)

            # 额外等待，确保蓝牙资源完全释放
            time.sleep(0.5)

        # 确保清理完成
        if not self._cleanup_done:
            # 清理未完成，再等待一下
            time.sleep(1.0)

        self.thread = None

    def get_data(self, timeout=1.0):
        """
        从队列获取IMU数据

        Returns:
            dict or None: IMU数据
        """
        try:
            return self.data_queue.get(timeout=timeout)
        except:
            return None

    def is_running(self):
        """检查是否正在运行"""
        return self.running and self.parse_enabled and self.thread and self.thread.is_alive()


if __name__ == "__main__":
    # 测试代码
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(levelname)s - %(message)s'
    )

    print("=" * 60)
    print("实时IMU监控测试")
    print("=" * 60)

    monitor = RealtimeIMUMonitor(mac_address="26:EF:60:96:87:3A")
    monitor.start()

    print("\n开始接收IMU数据...\n")

    while True:
        data = monitor.get_data(timeout=1.0)
        if data:
            print(f"\n📊 IMU数据: {data['timestamp']}")
            print(f"   加速度: ax={data['ax']:.3f}, ay={data['ay']:.3f}, az={data['az']:.3f}")
            print(f"   角速度: gx={data['gx']:.3f}, gy={data['gy']:.3f}, gz={data['gz']:.3f}")
            print(f"   角度: roll={data['roll']:.2f}, pitch={data['pitch']:.2f}, yaw={data['yaw']:.2f}")
            print(f"   振动幅度: {data['magnitude']:.3f}g")

        time.sleep(0.1)

    monitor.stop()
    print("测试完成")
